from .common import *

# Draw a broken horizontal bar chart.
def DrawBrokenBarHFig(x_values_set, y_values_set, legend_idx, legend_labels,
                      x_ticks, x_ticklabels, y_ticks, y_ticklabels,
                      x_label, y_label, x_min, x_max, y_min, y_max, bar_width,
                      folder, filename, allow_legend, allow_grid,
                      show_color=True, color_map = COLOR_MAP,
                      show_patterns=True, patterns = PATTERNS):
    """DrawBrokenBarHFig: Function used to draw broken horizontal bars.
    \t Usage:
    \t\t x_values_set: List of x values for all broken bars.
    \t\t y_values_set: List of y values for all broken bar sets.
    \t\t legend_idx: The index for the legend x values.
    \t\t legend_labels: Legend labels.
    \t\t x_ticks: Ticks along the x axis.
    \t\t x_ticklabels: Labels for each x tick. #x_ticklabels = #x_values_set[legend_idx]
    \t\t y_ticks: Ticks along the y axis.
    \t\t y_ticklabels: Labels for each broken bar set.
    \t\t {x,y}_label: Label for the x or y axis.
    \t\t x_{min,max}: The min and max value along the x axis.
    \t\t y_{min,max}: The min and max value along the y axis.
    \t\t bar_width: The vertical width of each broken bar.
    \t\t folder: The folder name.
    \t\t filename: The file name.
    \t\t allow_legend: The legend flag.
    \t\t allow_grid: The grid flag.
    \t\t show_color: Whether use color mode.
    \t\t color_map: The color map for colorful figures.
    \t\t show_patterns: Whether use pattern mode.
    \t\t patterns: The pattern palette.
    """
    assert(len(x_values_set) == len(y_values_set))
    assert(len(patterns) >= len(x_values_set))
    assert(len(x_ticks) == len(x_ticklabels))
    assert(len(y_ticks) == len(y_ticklabels))
    assert(len(color_map) >= len(x_values_set[0]))

    # figure size setting.
    fig = plt.figure(figsize=(8,3.2))
    figure = fig.add_subplot(111)

    # draw broken horizontal bars.
    broken_hbars = [None] * (len(y_values_set))

    for i in range(len(x_values_set)):
        if show_color and show_patterns:
            broken_hbars[i] = figure.broken_barh(x_values_set[i],
                                                 (y_values_set[i], bar_width),
                                                 edgecolor='black',
                                                 facecolors=color_map,
                                                 hatch=patterns[i])
        elif show_color:
            broken_hbars[i] = figure.broken_barh(x_values_set[i],
                                                 (y_values_set[i], bar_width),
                                                 edgecolor='black',
                                                 facecolors=color_map)
        else:
            broken_hbars[i] = figure.broken_barh(x_values_set[i],
                                                 (y_values_set[i], bar_width),
                                                 edgecolor='black',
                                                 hatch=patterns[i])

    # Add legend to horizontal bars.
    if allow_legend:
        assert(len(x_values_set[legend_idx]) == len(legend_labels))
        legend_patches = [None] * len(x_values_set[legend_idx])
        for i in range(len(legend_patches)):
            legend_patches[i] = mpatch.Rectangle((0, 0), 1, 1,
                                                 facecolor=color_map[i])

        figure.legend(legend_patches, legend_labels, prop=LEGEND_FP,
                      loc='upper center', ncol=len(legend_labels),
                      bbox_to_anchor=(0.45, 1.2), shadow=False,
                      frameon=False, borderpad=0.33, borderaxespad=0.0,
                      handlelength=2, labelspacing=0.2)

    # Set x_ticks and y_ticks.
    if x_ticks:
        plt.xticks(x_ticks, x_ticklabels)
    if y_ticks:
        plt.yticks(y_ticks, y_ticklabels)

    plt.xlim(x_min, x_max)
    plt.ylim(y_min, y_max)

    figure.grid(allow_grid, axis='both', color='gray', linestyle='--')
    figure.yaxis.set_major_locator(LinearLocator(6))

    figure.get_xaxis().set_tick_params(direction='in', pad=10)
    figure.get_yaxis().set_tick_params(direction='in', pad=10)

    plt.xlabel(x_label, fontproperties=LABEL_FP)
    plt.ylabel(y_label, fontproperties=LABEL_FP)

    if folder and filename:
        if not os.path.exists(folder):
            os.makedirs(folder)
        plt.savefig(folder + "/" + filename + ".eps", bbox_inches='tight',
                    format='eps')

# Draw colorful broken bars.
def DrawColorBrokenBarHFig(x_values_set, y_values_set, legend_idx, legend_labels,
                           x_ticks, x_ticklabels, y_ticks, y_ticklabels,
                           x_label, y_label, x_min, x_max, y_min, y_max,
                           bar_width, folder, filename, allow_legend, allow_grid,
                           color_map):
    DrawBrokenBarHFig(x_values_set, y_values_set, legend_idx, legend_labels,
                   x_ticks, x_ticklabels, y_ticks, y_ticklabels,
                   x_label, y_label, x_min, x_max, y_min, y_max,
                   bar_width, folder, filename, allow_legend, allow_grid,
                   True, color_map, False)

# Draw hatch-based broken bars.
def DrawHatchBrokenBarHFig(x_values_set, y_values_set, legend_idx, legend_labels,
                           x_ticks, x_ticklabels, y_ticks, y_ticklabels,
                           x_label, y_label, x_min, x_max, y_min, y_max,
                           bar_width, folder, filename, allow_legend, allow_grid,
                           patterns):
    DrawBrokenBarHFig(x_values_set, y_values_set, legend_idx, legend_labels,
                   x_ticks, x_ticklabels, y_ticks, y_ticklabels,
                   x_label, y_label, x_min, x_max, y_min, y_max,
                   bar_width, folder, filename, allow_legend, allow_grid,
                   show_color=False, show_pattern=True, patterns=patterns)
